package cp_zkp

import (
	"crypto/rand"
	"fmt"
	"math/big"
	"os"

	"github.com/linchao/zkp_auth/lib/config"
	"github.com/linchao/zkp_auth/lib/util"
)

// CPZKP 代表Chaum-Pedersen零知识证明系统
type CPZKP struct {
	// 预定义的系统参数
	pValue string
	qValue string
	gValue string
	hValue string
}

// CPZKPParams 表示ZKP协议的公共参数
type CPZKPParams struct {
	p, q, g, h *big.Int
}

// Prover 表示ZKP协议中的证明者
type Prover struct {
	x *big.Int // 秘密数x
}

// Verifier 表示ZKP协议中的验证者
type Verifier struct{}

// NewCPZKP 创建一个新的CPZKP实例
//
// 如果环境变量中设置了自定义参数，则使用环境变量的值
// 否则使用config中定义的默认参数
//
// 返回:
//   - CPZKP实例指针
//   - 如果过程中发生错误，返回error
func NewCPZKP() (*CPZKP, error) {
	// 从环境变量获取自定义参数，如果未设置则使用config中的默认值
	pValue := getEnvOrDefault("ZKP_P_VALUE", config.CPZKP_PARAM_P)
	qValue := getEnvOrDefault("ZKP_Q_VALUE", config.CPZKP_PARAM_Q)
	gValue := getEnvOrDefault("ZKP_G_VALUE", config.CPZKP_PARAM_G)
	hValue := getEnvOrDefault("ZKP_H_VALUE", config.CPZKP_PARAM_H)

	// 创建并返回CPZKP实例
	return &CPZKP{
		pValue: pValue,
		qValue: qValue,
		gValue: gValue,
		hValue: hValue,
	}, nil
}

// getEnvOrDefault 获取环境变量的值，如果未设置则返回默认值
func getEnvOrDefault(key, defaultValue string) string {
	value := os.Getenv(key)
	if value == "" {
		return defaultValue
	}
	return value
}

// InitCPZKPParams 使用CPZKP实例中的参数初始化CPZKPParams
//
// 返回:
//   - CPZKPParams实例指针
//   - 如果参数无效，返回error
func (c *CPZKP) InitCPZKPParams() (*CPZKPParams, error) {
	return InitCPZKPParams(c.pValue, c.qValue, c.gValue, c.hValue)
}

// 为CPZKPParams添加访问器方法，以便能访问内部字段
func (p *CPZKPParams) P() *big.Int {
	return p.p
}

func (p *CPZKPParams) Q() *big.Int {
	return p.q
}

func (p *CPZKPParams) G() *big.Int {
	return p.g
}

func (p *CPZKPParams) H() *big.Int {
	return p.h
}

// InitCPZKPParams 初始化Chaum-Pedersen ZKP协议系统参数
func InitCPZKPParams(pValue, qValue, gValue, hValue string) (*CPZKPParams, error) {
	// 使用util包的ParseBigInt函数解析字符串为大整数
	p, err := util.ParseBigInt(pValue, "p")
	if err != nil {
		return nil, fmt.Errorf("无效的p参数: %v", err)
	}

	q, err := util.ParseBigInt(qValue, "q")
	if err != nil {
		return nil, fmt.Errorf("无效的q参数: %v", err)
	}

	g, err := util.ParseBigInt(gValue, "g")
	if err != nil {
		return nil, fmt.Errorf("无效的g参数: %v", err)
	}

	h, err := util.ParseBigInt(hValue, "h")
	if err != nil {
		return nil, fmt.Errorf("无效的h参数: %v", err)
	}

	// 安全检查: 确保参数满足CP-ZKP协议要求
	// 1. p应该是大素数
	if !isProbablePrime(p) {
		return nil, fmt.Errorf("p必须是素数")
	}

	// 2. q应该是(p-1)的因子并且也是素数
	if !isProbablePrime(q) {
		return nil, fmt.Errorf("q必须是素数")
	}

	pMinus1 := new(big.Int).Sub(p, big.NewInt(1))
	if new(big.Int).Mod(pMinus1, q).Cmp(big.NewInt(0)) != 0 {
		return nil, fmt.Errorf("q必须是(p-1)的因子")
	}

	// 3. g, h的阶应该是q
	gPowQ := new(big.Int).Exp(g, q, p)
	if gPowQ.Cmp(big.NewInt(1)) != 0 {
		return nil, fmt.Errorf("g^q mod p 必须等于 1")
	}

	hPowQ := new(big.Int).Exp(h, q, p)
	if hPowQ.Cmp(big.NewInt(1)) != 0 {
		return nil, fmt.Errorf("h^q mod p 必须等于 1")
	}

	return &CPZKPParams{
		p: p,
		q: q,
		g: g,
		h: h,
	}, nil
}

// 检查一个数是否可能是素数
func isProbablePrime(n *big.Int) bool {
	// 使用Miller-Rabin算法进行素性测试，执行30轮测试
	return n.ProbablyPrime(30)
}

// NewProver 创建一个新的Prover，使用给定的秘密密码x
func NewProver(x *big.Int) *Prover {
	return &Prover{x: x}
}

// GenerateYValues 生成公钥值y1和y2
func (p *Prover) GenerateYValues(params *CPZKPParams) (y1, y2 *big.Int) {
	// 安全计算: 使用模幂运算计算公钥值
	// y1 = g^x mod p
	// y2 = h^x mod p
	y1 = new(big.Int).Exp(params.g, p.x, params.p)
	y2 = new(big.Int).Exp(params.h, p.x, params.p)
	return y1, y2
}

// CreateProofCommitment 创建ZKP承诺
func (p *Prover) CreateProofCommitment(params *CPZKPParams) (k, r1, r2 *big.Int, err error) {
	// 安全随机数: 生成0<k<q的随机值
	k, err = rand.Int(rand.Reader, params.q)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("生成随机数失败: %v", err)
	}

	// 安全检查: 确保k不为0
	if k.Cmp(big.NewInt(0)) <= 0 {
		return p.CreateProofCommitment(params) // 重新生成
	}

	// 计算承诺值
	// r1 = g^k mod p
	// r2 = h^k mod p
	r1 = new(big.Int).Exp(params.g, k, params.p)
	r2 = new(big.Int).Exp(params.h, k, params.p)

	return k, r1, r2, nil
}

// CreateProofChallenge 创建验证挑战
func (v *Verifier) CreateProofChallenge(params *CPZKPParams) (c *big.Int, err error) {
	// 安全随机数: 生成0<c<q的随机挑战值
	c, err = rand.Int(rand.Reader, params.q)
	if err != nil {
		return nil, fmt.Errorf("生成随机挑战失败: %v", err)
	}

	// 安全检查: 确保c不为0
	if c.Cmp(big.NewInt(0)) <= 0 {
		return v.CreateProofChallenge(params) // 重新生成
	}

	return c, nil
}

// CreateProofChallengeResponse 创建证明响应
func (p *Prover) CreateProofChallengeResponse(k, c *big.Int, params *CPZKPParams) (s *big.Int) {
	// 计算响应值 s = (k - c*x) mod q
	cx := new(big.Int).Mul(c, p.x)
	s = new(big.Int).Sub(k, cx)

	// 确保s为正值（处理模运算中的负数）
	if s.Sign() < 0 {
		s.Add(s, params.q)
	}
	s.Mod(s, params.q)

	return s
}

// VerifyProof 验证ZKP证明
func (v *Verifier) VerifyProof(y1, y2, r1, r2, c, s *big.Int, params *CPZKPParams) bool {
	// 安全检查: 验证所有输入参数不为nil
	if y1 == nil || y2 == nil || r1 == nil || r2 == nil || c == nil || s == nil || params == nil {
		return false
	}

	// 安全检查: 验证s, c在合法范围内(0 < s,c < q)
	if s.Cmp(big.NewInt(0)) <= 0 || s.Cmp(params.q) >= 0 ||
		c.Cmp(big.NewInt(0)) <= 0 || c.Cmp(params.q) >= 0 {
		return false
	}

	// 验证等式1: g^s * y1^c = r1 (mod p)
	// 即 g^s * (g^x)^c = g^k (mod p)
	gs := new(big.Int).Exp(params.g, s, params.p)
	y1c := new(big.Int).Exp(y1, c, params.p)
	l1 := new(big.Int).Mul(gs, y1c)
	l1.Mod(l1, params.p)

	if l1.Cmp(r1) != 0 {
		return false
	}

	// 验证等式2: h^s * y2^c = r2 (mod p)
	// 即 h^s * (h^x)^c = h^k (mod p)
	hs := new(big.Int).Exp(params.h, s, params.p)
	y2c := new(big.Int).Exp(y2, c, params.p)
	l2 := new(big.Int).Mul(hs, y2c)
	l2.Mod(l2, params.p)

	return l2.Cmp(r2) == 0
}
